from typing import List


class Solution:
    def removeBoxes(self, boxes: List[int]) -> int:
        # 定义状态矩阵：memo[i,j,k] = 从i到j中移除颜色相同的盒子，j右边有k个和a[j]颜色相同的盒子
        memo = {}

        def dp(l, r, k):
            print(l, r, k)
            if l > r:
                return 0
            if (l, r, k) in memo:
                return memo[(l, r, k)]
            while r > l and boxes[r] == boxes[r - 1]:
                r -= 1
                k += 1
            res = dp(l, r - 1, 0) + (k + 1) * (k + 1)

            for i in range(l, r):
                if boxes[i] == boxes[r]:
                    res = max(res, dp(i + 1, r - 1, 0) + dp(l, i, k + 1))

            memo[(l, r, k)] = res
            return res

        return dp(0, len(boxes) - 1, 0)


if __name__ == "__main__":
    print(Solution().removeBoxes([1, 3, 2, 2, 2, 3, 4, 3, 1]))  # 23

    print(Solution().removeBoxes([1, 1, 1]))  # 9
    print(Solution().removeBoxes([1]))  # 1

    # 测试用例20
    print(Solution().removeBoxes(
        [1, 2, 2, 1, 1, 1, 2, 1, 1, 2, 1, 2, 1, 1, 2, 2, 1, 1, 2, 2, 1, 1, 1, 2, 2, 2, 2, 1, 2, 1, 1, 2, 2, 1, 2, 1, 2,
         2, 2, 2, 2, 1, 2, 1, 2, 2, 1, 1, 1, 2, 2, 1, 2, 1, 2, 2, 1, 2, 1, 1, 1, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2, 2, 1, 1,
         1, 1, 1, 2, 2, 2, 2, 2, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 1, 1, 1, 2, 1, 2, 2, 1]))

    # 自制用例
    print(Solution().removeBoxes([1, 2, 3, 4, 5, 6, 1, 6, 5, 4, 3, 2, 1]))  # 25
